from utils.trie import Trie


class Gazetteer:
    def __init__(self, lower):
        """
        初始化 Gazetteer 实体词典
        :param lower: 是否进行大小写转换的标志。如果为 True，则所有单词都转换为小写。
        """
        self.trie = Trie()# 创建一个 Trie 数据结构用于存储和匹配词条
        self.ent2type = {}  ## word list to type 用于存储每个词条对应的实体类型（如人名、地名等）
        self.ent2id = {"<UNK>": 0}  ## word list to id # 用于存储每个词条对应的唯一 ID，初始化时添加未知词项
        self.lower = lower # 是否进行小写转换
        self.space = "" # 用于连接单词的分隔符，默认为空字符（即直接连接）

    def enumerateMatchList(self, word_list):
        """
       在词典中查找与输入词列表匹配的所有实体，是匹配所有词的所有前缀
       :param word_list: 输入的词列表，可能包含多个单词
       :return: 匹配的实体列表
       """
        if self.lower:
            word_list = [word.lower() for word in word_list]# 将所有单词转换为小写
        match_list = self.trie.enumerateMatch(word_list, self.space)# 在 Trie 中查找匹配的实体
        return match_list # 返回匹配的实体列表

    def insert(self, word_list, source):
        """
        向词典中插入新的实体词条
        :param word_list: 实体对应的词列表（一个实体可能是多个词组成的）这个应该是一个单词
        :param source: 实体的来源（例如实体类型或者其他标识）
        """
        if self.lower:
            word_list = [word.lower() for word in word_list]# 将所有单词转换为小写
        self.trie.insert(word_list) # 将词列表插入 Trie 中
        string = self.space.join(word_list) # 将词列表连接成一个字符串
        if string not in self.ent2type: # 如果该实体尚未在 ent2type 中记录
            self.ent2type[string] = source # 记录该实体的类型
        if string not in self.ent2id:# 如果该实体尚未在 ent2id 中记录
            self.ent2id[string] = len(self.ent2id) # 给该实体分配一个唯一的 ID

    def searchId(self, word_list):
        """
        查找输入词条对应的 ID
        :param word_list: 输入的词list ,一个单词
        :return: 对应实体的 ID，如果未找到，则返回默认的 <UNK> ID
        """
        if self.lower:
            word_list = [word.lower() for word in word_list]
        string = self.space.join(word_list) # 将词列表连接成一个字符串
        if string in self.ent2id:
            return self.ent2id[string]# 返回该实体的 ID
        return self.ent2id["<UNK>"] # 如果没有找到实体，则返回 <UNK> 对应的 ID

    def searchType(self, word_list):
        """
        查找输入词条对应的类型
        :param word_list: 输入的词列表
        :return: 对应实体的类型，如果未找到，则抛出错误并退出程序
        """
        if self.lower:
            word_list = [word.lower() for word in word_list]
        string = self.space.join(word_list) #将词列表连接成一个字符串
        if string in self.ent2type: # 如果该实体存在于 ent2type 中
            return self.ent2type[string] # 返回该实体的类型
        print(
            "Error in finding entity type at gazetteer.py, exit program! String:",
            string,
        )
        exit(0)

    def size(self):
        """
        获取词典中已记录的实体数量
        :return: 词典中实体的数量
        """
        return len(self.ent2type)
